#!/usr/bin/env python3

import argparse
import pendulum
import os
import sys
script_root = os.path.dirname(os.path.realpath(__file__))
sys.path.append(f'{script_root}/../utils')
from utils import cli, check_files, edit_file, has_key, run, submit_job

def run_upp(work_root, upp_root, config, args, tag=None, time=None):
	start_time = config['custom']['start_time']
	end_time = config['custom']['end_time']

	if time != None:
		times = [time]
	elif args.time:
		times = [args.time]
	else:
		times = []
		if has_key(config['custom'], ['upp', 'wrfout_interval_seconds']):
			wrfout_interval_seconds = config['custom']['upp']['wrfout_interval_seconds']
		else:
			wrfout_interval_seconds = 3600
		for time in pendulum.period(start_time, end_time).range('seconds', wrfout_interval_seconds):
			times.append(time)

	if has_key(config['custom'], ['upp', 'doms']):
		doms = config['custom']['upp']['doms']
	else:
		doms = ['d01']

	if tag:
		wrf_work_dir = f'{work_root}/{tag}/wrf'
	else:
		wrf_work_dir = f'{work_root}/wrf'
	if not os.path.isdir(wrf_work_dir): cli.error('WRF is not run successfully!')

	if tag:
		upp_work_dir = f'{work_root}/{tag}/upp'
	else:
		upp_work_dir = f'{work_root}/upp'
	if not os.path.isdir(upp_work_dir): os.makedirs(upp_work_dir)

	if has_key(config['custom'], ['upp', 'postxconfig_txt']):
		run(f'ln -sf {os.path.abspath(config["custom"]["upp"]["postxconfig_txt"])} {upp_work_dir}/postxconfig-NT.txt')
	else:
		run(f'ln -sf {upp_root}/parm/postxconfig-NT-WRF.txt {upp_work_dir}/postxconfig-NT.txt')
	run(f'ln -sf {upp_root}/parm/post_avblflds.xml {upp_work_dir}/post_avblflds.xml')
	run(f'ln -sf {upp_root}/parm/params_grib2_tbl_new {upp_work_dir}/params_grib2_tbl_new')
	run(f'ln -sf {upp_root}/parm/nam_micro_lookup.dat {upp_work_dir}/nam_micro_lookup.dat')
	run(f'ln -sf {upp_root}/parm/hires_micro_lookup.dat {upp_work_dir}/hires_micro_lookup.dat')
	crtm_dir = f'{upp_root}/sorc/comlibs/crtm2/src/fix'
	run(f'ln -sf {crtm_dir}/EmisCoeff/IR_Water/Big_Endian/Nalli.IRwater.EmisCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/EmisCoeff/MW_Water/Big_Endian/FASTEM4.MWwater.EmisCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/EmisCoeff/MW_Water/Big_Endian/FASTEM5.MWwater.EmisCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/EmisCoeff/MW_Water/Big_Endian/FASTEM6.MWwater.EmisCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/EmisCoeff/IR_Land/SEcategory/Big_Endian/NPOESS.IRland.EmisCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/EmisCoeff/IR_Snow/SEcategory/Big_Endian/NPOESS.IRsnow.EmisCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/EmisCoeff/IR_Ice/SEcategory/Big_Endian/NPOESS.IRice.EmisCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/AerosolCoeff/Big_Endian/AerosolCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/CloudCoeff/Big_Endian/CloudCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/imgr_g11.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/imgr_g11.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/imgr_g12.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/imgr_g12.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/imgr_g13.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/imgr_g13.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/imgr_g15.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/imgr_g15.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/imgr_mt1r.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/imgr_mt1r.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/imgr_mt2.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/imgr_mt2.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/imgr_insat3d.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/imgr_insat3d.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/amsre_aqua.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/amsre_aqua.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/tmi_trmm.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/tmi_trmm.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/ssmi_f13.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/ssmi_f13.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/ssmi_f14.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/ssmi_f14.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/ssmi_f15.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/ssmi_f15.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/ssmis_f16.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/ssmis_f16.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/ssmis_f17.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/ssmis_f17.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/ssmis_f18.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/ssmis_f18.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/ssmis_f19.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/ssmis_f19.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/ssmis_f20.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/ssmis_f20.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/seviri_m10.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/seviri_m10.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/v.seviri_m10.SpcCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/TauCoeff/ODPS/Big_Endian/abi_gr.TauCoeff.bin {upp_work_dir}/')
	run(f'ln -sf {crtm_dir}/SpcCoeff/Big_Endian/abi_gr.SpcCoeff.bin {upp_work_dir}/')

	old_dir = os.getcwd()
	os.chdir(upp_work_dir)
	for time in times:
		for dom in doms:
			fhr = int((time - start_time).total_seconds() / 3600)
			output_file = f'WRFPRS.{dom}.GrbF{fhr:02}'
			if not os.path.isfile(output_file):
				wrfout_file = f'{wrf_work_dir}/wrfout_{dom}_{time.format("YYYY-MM-DD_HH:mm:SS")}'
				if not os.path.isfile(wrfout_file):
					cli.warning(f'WRF output file {wrfout_file} does not exist! Skip it ...')
					continue
				f = open('itag', 'w')
				f.write(f'{wrfout_file}\n')
				f.write('netcdf\n')
				f.write('grib2\n')
				f.write(time.format('YYYY-MM-DD_HH:mm:SS')+'\n')
				f.write('NCAR\n')
				f.close()
				for _ in range(10):
					submit_job(f'{upp_root}/exec/unipost.exe', 4, config, args, wait=True)
					if os.path.isfile(f'WRFPRS.GrbF{fhr:02}'):
						run(f'mv WRFPRS.GrbF{fhr:02} {output_file}')
						break
					else:
						cli.warning(f'Failed to run UPP at {time}! Try again!')
				if not os.path.isfile(output_file):
					cli.error(f'Failed to run UPP at {time}!')
	run(f'rm *.bin *.dat *.xml fort.110 *.txt itag params_grib2_tbl_new')
	os.chdir(old_dir)
